package com.lw.leetcode.sort.c;

import com.lw.test.util.Utils;

import java.util.PriorityQueue;

/**
 * Created with IntelliJ IDEA.
 * 1439. 有序矩阵中的第 k 个最小数组和
 *
 * @author liw
 * @version 1.0
 * @date 2022/6/28 13:41
 */
public class KthSmallest {

    public static void main(String[] args) {
        KthSmallest test = new KthSmallest();

        // 7
//        int[][] arr = {{1, 3, 11}, {2, 4, 6}};
//        int k = 5;

        // 17
//        int[][] arr = {{1, 3, 11}, {2, 4, 6}};
//        int k = 9;

        // 9
//        int[][] arr = {{1, 10, 10}, {1, 4, 5}, {2, 3, 6}};
//        int k = 7;

        // 12
//        int[][] arr = {{1, 1, 10}, {2, 2, 9}};
//        int k = 7;

        //
        int[][] arr = Utils.getArr(40, 40, 1, 1);
        int k = 200;
        System.out.println(k);


        int i = test.kthSmallest(arr, k);
        System.out.println(i);
    }

    public int kthSmallest(int[][] mat, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>((a, b) -> Integer.compare(b, a));
        int m = mat.length;
        int[] ms = mat[0];
        int n = ms.length;
        int l = Math.min(n, k);
        for (int i = 0; i < l; i++) {
            queue.add(ms[i]);
        }
        int[] arr = new int[k];
        for (int i = 1; i < m; i++) {
            ms = mat[i];
            int size = queue.size();
            for (int index = size - 1; index >= 0; index--) {
                arr[index] = queue.poll();
            }
            for (int v : ms) {
                if (queue.size() == k && v + arr[0] > queue.peek()) {
                    break;
                }
                for (int j = 0; j < size; j++) {
                    int t = arr[j] + v;
                    if (queue.size() < k) {
                        queue.add(t);
                        continue;
                    }
                    if (t > queue.peek()) {
                        break;
                    }
                    queue.poll();
                    queue.add(t);
                }
            }
        }
        return queue.poll();
    }

    public int kthSmallest2(int[][] mat, int k) {
        PriorityQueue<Integer> queue = new PriorityQueue<>();
        int m = mat.length;
        int n = mat[0].length;
        int sum = 0;
        for (int i = 0; i < m; i++) {
            int v = mat[i][0];
            sum += v;
            queue.add((v << 16) + (i << 8));
        }
        while (k > 1) {
            int poll = queue.poll();
            int v = poll >> 16;
            int x = (poll >> 8) & 0XFF;
            int y = poll & 0XFF;
            int t = Integer.MAX_VALUE;
            if (y + 1 < n) {
                t = mat[x][y + 1];
                queue.add((t << 16) + (x << 8) + y + 1);
            }
            while (!queue.isEmpty()) {
                int peek = queue.peek();
                if ((peek >> 16) == v) {
                    x = (poll >> 8) & 0XFF;
                    y = poll & 0XFF;
                    if (y + 1 < n) {
                        t = Math.min(t, mat[x][y + 1]);
                        queue.add((t << 16) + (x << 8) + y + 1);
                    }
                } else {
                    break;
                }
            }
            sum = sum - v + t;
            k--;
        }
        return sum;
    }

}
